package com.fudan.run.trans;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

import org.apache.spark.api.java.JavaRDD;

import com.fudan.cfg.base.TransDef;
import com.fudan.cfg.trans.CombineNode;
import com.fudan.cfg.trans.TargetNode;
import com.fudan.run.JobRunner;
import com.fudan.run.VariableRddMap;
import com.fudan.run.ctx.annotation.TransAdapter;

import scala.Tuple2;

/**
 * {
			"name": "after_combine_name",
			"type": "combine",
			"target": [
				{
					"name": "rdd_1",
					"role": "node",
					"as": "node1",
					"by":"id_1"
				},
				{
					"name": "rdd_2",
					"role": "node",
					"as": "node2",
					"by":"id_2"
				}
			]
		}
 * @author admin
 *
 */
@TransAdapter("combine")
public class CombineTrans extends TransHandler<CombineNode>{

	@Override
	public void trans(JobRunner jobRunner,CombineNode transNode, VariableRddMap variableMap) {
		List<TargetNode> targetNodeList = transNode.getTargetNodeList();
		if(targetNodeList.size() != 2) {
			throw new RuntimeException("combine节点的target数量异常");
		}
		TargetNode targetNode1 = targetNodeList.get(0);
		TargetNode targetNode2 = targetNodeList.get(1);
		String role1 = targetNode1.getRole();
		String role2 = targetNode2.getRole();
		//node,node
		if(role1.equals(role2) && "node".equals(role1)) {
			JavaRDD<Map<String, Object>> after = variableMap.getRdd(targetNode1.getName())
			.mapToPair(l->Tuple2.apply(getOneValueFromMapByKey(l,targetNode1.getBy()),l))
			.cogroup(variableMap.getRdd(targetNode2.getName())
					.mapToPair(l->Tuple2.apply(getOneValueFromMapByKey(l,targetNode2.getBy()),l)),transNode.getPartition())
			.map(line->{
				Map<String,Object> map = new HashMap<>();
				map.put(targetNode1.getAs(), new ArrayList<>());
				map.put(targetNode2.getAs(), new ArrayList<>());
				Iterator<Map<String, Object>> group1 = line._2._1.iterator();
				Iterator<Map<String, Object>> group2 = line._2._2.iterator();
				Object key = line._1;
				while(group1.hasNext()) {
					Map<String, Object> item = group1.next();
					((List<Map<String, Object>>)map.get(targetNode1.getAs())).add(item);
				}
				while(group2.hasNext()) {
					Map<String, Object> item = group2.next();
					((List<Map<String, Object>>)map.get(targetNode2.getAs())).add(item);
				}
				map.put(transNode.getKeyAs(), key);
				return map;
			});
			variableMap.addRdd(transNode.getName(), after);

		}
		//base,node
		else {
			if("base".equals(role1)) {
				JavaRDD<Map<String, Object>> after = variableMap.getRdd(targetNode1.getName())
						.mapToPair(l->Tuple2.apply(getOneValueFromMapByKey(l,targetNode1.getBy()),l))
						.cogroup(variableMap.getRdd(targetNode2.getName())
								.mapToPair(l->Tuple2.apply(getOneValueFromMapByKey(l,targetNode2.getBy()),l)),transNode.getPartition())
						.flatMap(line->{
							List<Map<String,Object>> ret = new ArrayList<>();
							Iterator<Map<String, Object>> group1 = line._2._1.iterator();
							Iterator<Map<String, Object>> group2 = line._2._2.iterator();
							List<Map<String,Object>> nodeValues = new ArrayList<>();
							if(!group1.hasNext()) {
								return ret.iterator();
							}
							while(group2.hasNext()) {
								Map<String, Object> item2 = group2.next();
								nodeValues.add(item2);
							}
							while(group1.hasNext()) {
								Map<String, Object> item = group1.next();
								item.put(targetNode2.getAs(), nodeValues);
								ret.add(item);
							}
							return ret.iterator();
						}).filter(f->f.size()>0);
						variableMap.addRdd(transNode.getName(), after);
			}else if("base".equals(role2)) {
				JavaRDD<Map<String, Object>> after = variableMap.getRdd(targetNode1.getName())
						.mapToPair(l->Tuple2.apply(getOneValueFromMapByKey(l,targetNode1.getBy()),l))
						.cogroup(variableMap.getRdd(targetNode2.getName())
								.mapToPair(l->Tuple2.apply(getOneValueFromMapByKey(l,targetNode2.getBy()),l)),transNode.getPartition())
						.flatMap(line->{
							List<Map<String,Object>> ret = new ArrayList<>();
							Iterator<Map<String, Object>> group1 = line._2._1.iterator();
							Iterator<Map<String, Object>> group2 = line._2._2.iterator();
							List<Map<String,Object>> nodeValues = new ArrayList<>();
							if(!group2.hasNext()) {
								return ret.iterator();
							}
							while(group1.hasNext()) {
								Map<String, Object> item1 = group1.next();
								nodeValues.add(item1);
							}
							while(group2.hasNext()) {
								Map<String, Object> item = group2.next();
								item.put(targetNode1.getAs(), nodeValues);
								ret.add(item);
							}
							return ret.iterator();
						}).filter(m->m.size()>0);
						variableMap.addRdd(transNode.getName(), after);
			}else {
				throw new RuntimeException("combine节点没有找到base角色");
			}
		}
	}

	private static Object getOneValueFromMapByKey(Map<String,Object> map,String key) {
		Object val = null;
		String[] keys = key.split("[.]");
		for (String name : keys) {
			if(val == null) {
				if(name.endsWith("[]")) {
					name = name.substring(0,name.indexOf("["));
					val = ((List<Map<String,Object>>)map.get(name)).get(0);
				}else {
					val = map.get(name);
				}
			}else {
				if(name.endsWith("[]")) {
					name = name.substring(0,name.indexOf("["));
					val = ((List<Map<String,Object>>)((Map<String,Object>)val).get(name)).get(0);
				}else {
					val = ((Map<String,Object>)val).get(name);
				}
			}
		}
		return val;
	}
}
